"""
A class to exercise the JT core functions.
"""
import itk
import vtk

import connect2
import ImageCompare

import JT
import Calibration

def Test():
    """
    Test out the ITK registration procedure.
    """

    inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
    staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
    calFileLat = 'C:/Users/bryan/bryan-code/trunk/Tests/ext_cal1.txt'
    #drrImage = None

    # These should be normalized images
    fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mha'
    projectedImageName = 'C:/Users/bryan/bryan-code/trunk/Tests/ProjectedImage.mhd'
    image1Name = 'C:/Users/bryan/bryan-code/trunk/Tests/Test_ImageOverlayFixed.mhd'
    image2Name = 'C:/Users/bryan/bryan-code/trunk/Tests/Test_ImageOverlayMoving.mhd'
    outPath = 'C:/Users/bryan/bryan-code/trunk/Tests/'

    #
    #  ITK Image Components
    #%#####################
    itk_reader = itk.ImageFileReader.IF2.New()
    itk_reader.SetFileName(fixedImageName)
    itk_reader.Update()

    # Create Registraion object
    reg = JT.Registration()

    #
    # Create Lateral DRR
    #%###########################
    xraycamLat = JT.XRayCamera()
    drrLat = JT.DRR()
    calLat = Calibration.ExternalCalibration()

    calLat.LoadConsolidatedCalibration(calFileLat)
    calLat.LoadStaFile(staFile)
    drrLat.SetXRayCamera(xraycamLat)
    drrLat.SetVolumeFilename(inputVolFileName)
    drrLat.SetBlackOnWhiteImage()

    #
    # Automatic calibration
    #%##########################
    xraycamLat.SetExternalCalibration(calLat)

    vtkMatrixLAT = calLat.ConvertNumpyToVTKArray(calLat._VolumeT)
    drrLat.SetVolumeUserMatrix(vtkMatrixLAT)
    #drrLat.Render()

    # Set the DRR projector to the registration
    reg.SetDRR(drrLat)
    reg.SetFixedImage(itk_reader.GetOutput())
    reg.SetupITKRegistration()
    reg.StartRegistration()

    print "Metric Value: %f" % reg.GetMetricValue()

    fixedImage = itk_reader.GetOutput()
    movingImage = drrLat.GetImage('itk')

    # Now, we use the final transform for resampling the moving image.
    resampler = itk.ResampleImageFilter.IF2IF2.New()

    resampler.SetTransform( reg._transform.GetPointer() )
    resampler.SetInput( movingImage )

    region = fixedImage.GetLargestPossibleRegion()

    resampler.SetSize( region.GetSize() )
    resampler.SetOutputSpacing( fixedImage.GetSpacing() )
    resampler.SetOutputDirection( fixedImage.GetDirection() )
    resampler.SetOutputOrigin(  fixedImage.GetOrigin() )
    resampler.SetDefaultPixelValue( 100 )

    #
    # Cast for output
    #
    outputCast = itk.RescaleIntensityImageFilter.IF2IUC2.New()
    #outputCast.SetInput( resampler.GetOutput() )
    outputCast.SetInput( movingImage )
    outputCast.SetOutputMinimum( 0 )
    outputCast.SetOutputMaximum( 255 )

    writer = itk.ImageFileWriter.IUC2.New()

    writer.SetFileName( 'C:/Users/bryan/bryan-code/trunk/Tests/output_registration.png' )
    writer.SetInput( outputCast.GetOutput() )
    writer.Update()
    writer.Write()


if __name__ == "__main__":
    Test()